# 抽象的DFT矩阵分解基类
import pickle
import numpy as np
from solutions.problem_2 import DFT_2
import os
from solvers import Individual_3, AlphaInherit


class DFT_3(DFT_2):
    @property
    def __name__(self):
        return "Problem_3"

    def train(
        self, lr: float = 0.01, epoch: int = 1000, store_path="./result/Problem_3.pkl", is_post_train: bool = False, **kwargs
    ):
        error = 100000
        error_list = []
        A_K_list = []
        beta_list = []
        self.lr=lr
        for i in range(epoch):
            try:
                A_K_temp = [
                    self.identify(x - lr * self.grad(k))
                    for k, x in enumerate(self._A_K)
                ]
                beta_temp = self.beta - lr * self.grad_beta()
                if self.error - error > 0:
                    break
                if i % 100 == 0:
                    # error = self.error
                    print(f"epoch:{i}, error:{self.error}")
                error = self.error
                self._A_K = A_K_temp
                self.beta = beta_temp
                error_list.append(self.error)
                A_K_list.append(self._A_K)
                beta_list.append(self.beta)
            except:
                raise Exception("Here! Error!")
        pickle.dump((self, error_list, A_K_list, beta_list), open(os.path.join(
                    store_path, f"K{self.K}_N{self.N}_lr{lr}_"+self.__name__+".pkl"), "wb"))
        if is_post_train:
            self.post_train(store_path=store_path, **kwargs)

    def post_train(
        self,
        store_path: str = "./result/Problem_3.pkl",
        iter_num: int = 10000,
        individual_num: int = 8,
        offstring_num: int = 8,
    ):
        """
        后期操作，在完成浮点训练之后，进行一次规划，由于规划的目标函数不可以求导，不能使用梯度下降，替代梯度的方法也不好做，这里尝试使用智能优化算法遗传算法来做。
        """
        print("*"*50)
        print("现在进入遗传算法内")
        print("*"*50)
        M, _, _, _ = pickle.load(
            open(os.path.join(store_path, f"K{self.K}_N{self.N}_lr{self.lr}_"+self.__name__+".pkl"), "rb"))
        IS = AlphaInherit(
            individual_num,
            Individual_3,
            val=np.random.normal(35000, 400, individual_num),
            sigma=np.random.normal(8000, 200, individual_num),
            dft=M.dft,
            beta=M.get_beta,
            A_K=M.A_K,
            K=M.K,
            gate=-30,
            lr=1e-4
        )
        final_group = IS.train(
            epoch_num=iter_num,
            multiply_kwargs={"offstring_num": offstring_num},
            filter_kwargs={},
        )
        val = 0
        min_obj = None
        for i in final_group:
            if val > i.score(M.dft, M.A_K, M.K, M.get_beta):
                val = i.score(M.dft, M.A_K, M.K, M.get_beta)
                min_obj = i

        self._A_K = [self.convert_to_power_of_two(
            self.max_two(min_obj.val*A_K_i)) for A_K_i in M._A_K]
        self.alpha = min_obj.val
        print(
            f"拟合结果与真实值的差: {val}, 拟合因子: {min_obj.val}, 拟合因子的标准差: {min_obj.sigma}")
        pickle.dump(
            (self, min_obj),
            open(os.path.join(store_path, self.__name__ +
                 "_processed" + ".pkl"), "wb"),
        )
